# -*- coding: utf-8 -*-
'''
@File: 			cosyvoice.py
@Version: 		1.0
@Author: 		kztxt
@Email: 		kztxt@163.com
@CreatedAt: 	2025-09-30 10:44:08
'''

# here put the import lib
import io, os
import torchaudio
import numpy as np
from loguru import logger
from fastapi import APIRouter, UploadFile, Form, File
from fastapi.responses import StreamingResponse, Response
from cosyvoice2.cosyvoice.utils.file_utils import load_wav
from cosyvoice2.model import cosyvoice2
from pydub import AudioSegment

router = APIRouter()


def generate_data(model_output):
    for i in model_output:
        tts_audio = (i['tts_speech'].numpy() * (2 ** 15)).astype(np.int16).tobytes()
        yield tts_audio

@router.get("/inference_sft")
@router.post("/inference_sft")
async def inference_sft(tts_text: str = Form(), spk_id: str = Form()):
    model_output = cosyvoice2.inference_sft(tts_text, spk_id)
    return StreamingResponse(generate_data(model_output))


@router.get("/inference_zero_shot")
@router.post("/inference_zero_shot")
async def inference_zero_shot(tts_text: str = Form(), prompt_text: str = Form(), prompt_wav: UploadFile = File()):
    prompt_speech_16k = load_wav(prompt_wav.file, 16000)
    audio_buffer = None
    for i, j in enumerate(cosyvoice2.inference_zero_shot(tts_text, prompt_text, prompt_speech_16k, stream=False, speed=1.0)):
        buf = io.BytesIO()
        torchaudio.save(buf, j['tts_speech'], cosyvoice2.sample_rate, format='wav')
        audio = AudioSegment.from_wav(buf)
        if audio_buffer is None:
            audio_buffer = audio
        else:
            audio_buffer += audio
    wav_bs = io.BytesIO()
    audio_buffer.export(wav_bs, format='mp3')
    return Response(wav_bs.getvalue(), media_type="audio/mp3")


@router.get("/inference_cross_lingual")
@router.post("/inference_cross_lingual")
async def inference_cross_lingual(tts_text: str = Form(), prompt_wav: UploadFile = File()):
    prompt_speech_16k = load_wav(prompt_wav.file, 16000)
    model_output = cosyvoice2.inference_cross_lingual(tts_text, prompt_speech_16k)
    return StreamingResponse(generate_data(model_output))


@router.get("/inference_instruct")
@router.post("/inference_instruct")
async def inference_instruct(tts_text: str = Form(), spk_id: str = Form(), instruct_text: str = Form()):
    model_output = cosyvoice2.inference_instruct(tts_text, spk_id, instruct_text)
    return StreamingResponse(generate_data(model_output))


@router.get("/inference_instruct2")
@router.post("/inference_instruct2")
async def inference_instruct2(tts_text: str = Form(), instruct_text: str = Form(), prompt_wav: UploadFile = File()):
    prompt_speech_16k = load_wav(prompt_wav.file, 16000)
    model_output = cosyvoice2.inference_instruct2(tts_text, instruct_text, prompt_speech_16k)
    return StreamingResponse(generate_data(model_output))